/*=============================================================================
*
*   Copyright (C) 2020 All rights reserved.
*
*   Filename:		biot.hpp
*
*   Author: Wang Zhecheng - wangzhecheng@yeah.net
*
*   Date: 2020-03-04 17:57
*
*   Last Editors: Wang Zhecheng - wangzhecheng@yeah.net
*
*   Last modified:	2020-08-04 01:21
*
*   Description:
*
=============================================================================*/
#ifndef BIOT_H
#define BIOT_H

#include "io/vtu_writer.hpp"
#include "model/couple/displacement.hpp"
#include "model/couple/pore.hpp"
#include "model/solve/assemble_elemental_mat.hpp"
#include "model/time_step.hpp"

namespace geoxfem {

namespace model {
struct biot {};
}  // namespace model

namespace traits {
template <>
struct if_consider_deform<model::biot> {
  static constexpr bool value = true;
};
}  // namespace traits

namespace geometry {

template <>
template <typename Vertex>
std::vector<model::Dof*> get_nodal_dofs<model::biot> /*{{{*/
  ::apply(Vertex& vertex) {
  auto dof_list1 = get_nodal_dofs<model::displacement>::apply(vertex);
  auto dof_list2 = get_nodal_dofs<model::pore>::apply(vertex);
  dof_list1.insert(dof_list1.end(), dof_list2.begin(), dof_list2.end());
  return dof_list1;
} /*}}}*/

template <>
template <typename Vertex>
void set_nodal_dofs<model::biot>::apply(Vertex& v) { /*{{{*/
  set_nodal_dofs<model::displacement>::apply(v);
  set_nodal_dofs<model::pore>::apply(v);
  return;
} /*}}}*/

template <>
template <typename DimensionType, typename degree>
std::vector<model::Dof*> get_elemental_dof<model::biot>::/*{{{*/
  apply(const geometry::ShapeInterface<DimensionType, degree>& shp) {
  std::vector<model::Dof*> doflist1 =
    get_elemental_dof<model::displacement>::apply(shp);
  std::vector<model::Dof*> doflist2 =
    get_elemental_dof<model::pore>::apply(shp);
  doflist1.insert(doflist1.end(), doflist2.begin(), doflist2.end());
  return doflist1;
} /*}}}*/
}  // namespace geometry

namespace model {

template <>
template <typename DimensionType, typename degree>
MatrixXR assemble_gauss_kay<pore, displacement>::/*{{{*/
  apply(const geometry::ShapeInterface<DimensionType, degree>& shp,
        const GaussPoint<DimensionType>&                       gp) {
  static constexpr size_t dimension = traits::dimension<DimensionType>::value;

  VectorXR N   = ShapeFunction(shp, gp);
  MatrixXR bee = BeeMatrix(shp, gp);
  MatrixXR Mee;
  Mee.setZero(N.size(), bee.rows());
  Mee.leftCols(dimension).colwise() = N;
  /// Mee = N1, N1, N1, 0, 0, 0,
  ///       N2, N2, N2, 0, 0, 0,
  ///       N3, N3, N3, 0, 0, 0,
  ///       N4, N4, N4, 0, 0, 0,
  ///       ...
  return Mee * bee;
} /*}}}*/

template <>
template <typename DimensionType, typename degree, typename TimeStep>
MatrixXR assemble_elemental_kay<biot>::/*{{{*/
  apply(const geometry::ShapeInterface<DimensionType, degree>& shp,
        const TimeStep&                                        tstep) {
  size_t   neq = geometry::get_elemental_dof<biot>::apply(shp).size();
  MatrixXR ekay(neq, neq);
  ekay.setZero();

  MatrixXR ks = assemble_sub_kay<displacement, displacement>::apply(shp);
  MatrixXR kc = assemble_sub_kay<pore, displacement>::apply(shp);
  MatrixXR kp = assemble_sub_kay<pore, pore>::apply(shp);

  _real theta = time_theta::set_get();
  _real dtime = traits::get_delta_time::apply(tstep);

  ekay << ks, kc.transpose(),  //
    kc, -dtime * theta * kp;

  return ekay;
} /*}}}*/

template <>
template <typename DimensionType, typename degree>
VectorXR assemble_sub_rhs_for_tstep<pore, displacement> /*{{{*/
  ::apply(const geometry::ShapeInterface<DimensionType, degree>& shp) {
  size_t neq = geometry::get_elemental_dof<pore>::apply(shp).size();
  return VectorXR::Zero(neq);
} /*}}}*/

template <>
template <typename DimensionType, typename degree, typename TimeStep>
VectorXR assemble_element_rhs_for_tstep<biot> /*{{{*/
  ::apply(const geometry::ShapeInterface<DimensionType, degree>& shp,
          const TimeStep&                                        tstep) {

  size_t   neq   = geometry::get_elemental_dof<biot>::apply(shp).size();
  _real    dtime = tstep.getDurationTime();
  VectorXR eload11 =
    assemble_sub_rhs_for_tstep<displacement, displacement>::apply(shp);
  VectorXR eload12 = assemble_sub_rhs_for_tstep<displacement, pore>::apply(shp);
  VectorXR eload21 = assemble_sub_rhs_for_tstep<pore, displacement>::apply(shp);
  VectorXR eload22 =
    -dtime * assemble_sub_rhs_for_tstep<pore, pore>::apply(shp);
  GEOXFEM_ASSERT(eload11.rows() == eload12.rows());
  GEOXFEM_ASSERT(eload21.rows() == eload22.rows());
  GEOXFEM_ASSERT(eload11.rows() + eload22.rows() == neq);

  VectorXR eload(neq);
  eload << eload11 + eload12, eload21 + eload22;

  return eload;
} /*}}}*/

template <>
template <typename DimensionType, typename degree, typename TimeStep>
VectorXR assemble_element_rhs_for_iter<biot>::/*{{{*/
  apply(const geometry::ShapeInterface<DimensionType, degree>& shp,
        const TimeStep&                                        tstep) {

  size_t   neq   = geometry::get_elemental_dof<biot>::apply(shp).size();
  _real    dtime = tstep.getDurationTime();
  _real    theta = time_theta::set_get();
  VectorXR eload11 =
    assemble_sub_rhs_for_iter<displacement, displacement>::apply(shp);
  VectorXR eload12 = assemble_sub_rhs_for_iter<displacement, pore>::apply(shp);
  VectorXR eload21 = assemble_sub_rhs_for_iter<pore, displacement>::apply(shp);
  VectorXR eload22 =
    -dtime * theta * assemble_sub_rhs_for_iter<pore, pore>::apply(shp);
  GEOXFEM_ASSERT(eload11.rows() == eload12.rows());
  GEOXFEM_ASSERT(eload21.rows() == eload22.rows());
  GEOXFEM_ASSERT(eload11.rows() + eload22.rows() == neq);

  VectorXR eload(neq);
  eload << eload11 + eload12, eload21 + eload22;

  return eload;
} /*}}}*/

template <>
template <typename Vertex, typename TimeStep>
void update_nodal_boundary_condition<biot> /*{{{*/
  ::apply(const TimeStep& tstep, Vertex& vertex) {
  update_nodal_boundary_condition<displacement>::apply(tstep, vertex);
  update_nodal_boundary_condition<pore>::apply(tstep, vertex);
  return;
} /*}}}*/

}  // namespace model

namespace io {

template <>
template <typename Mesh>
void write_ptree<model::biot>::/*{{{*/
  apply(const Mesh& mesh, boost::property_tree::ptree& root) {
  using ptree      = boost::property_tree::ptree;
  ptree& PointData = root.get_child("VTKFile.UnstructuredGrid.Piece.PointData");
  ptree& CellData  = root.get_child("VTKFile.UnstructuredGrid.Piece.CellData");
  ////////////////////////////////////////////////////////////////
  // populate point data

  ptree deform = io::vtu_node_info::apply(mesh, io::vtu_vertex_deform{},
                                          {{"type", "Float32"},
                                           {"NumberOfComponents", "3"},
                                           {"Name", "displacement"}});
  PointData.add_child("DataArray", deform);

  ptree pore = io::vtu_node_info::apply(
    mesh, io::vtu_vertex_pore{}, {{"type", "Float32"}, {"Name", "pore"}});
  PointData.add_child("DataArray", pore);

  ptree v_xfem = io::vtu_node_info::apply(
    mesh, [](const auto& v) { return v.ifXfem() ? "1\n" : "0\n"; },
    {{"type", "Float32"}, {"Name", "v_xfem"}});
  PointData.add_child("DataArray", v_xfem);
  ////////////////////////////////////////////////////////////////
  // populate cell data

  ptree sigma = io::vtu_cell_info::apply(
    mesh, io::vtu_cell_sigma{},
    {{"type", "Float32"}, {"NumberOfComponents", "9"}, {"Name", "sigma"}});
  CellData.add_child("DataArray", sigma);

  ptree epsilon = io::vtu_cell_info::apply(
    mesh, io::vtu_cell_epsilon{},
    {{"type", "Float32"}, {"NumberOfComponents", "9"}, {"Name", "epsilon"}});
  CellData.add_child("DataArray", epsilon);

  ptree cell_xfem = io::vtu_node_info::apply(
    mesh, [](const auto& v) { return v.ifXfem() ? "1\n" : "0\n"; },
    {{"type", "Float32"}, {"Name", "cell_xfem"}});
  PointData.add_child("DataArray", cell_xfem);
  return;
} /*}}}*/

}  // namespace io
}  // namespace geoxfem
#endif /* BIOT_H */
